package middleware

import (
	"bytes"
	"com.yy.api/my_gateway/internal/cache"
	"fmt"
	"github.com/gin-gonic/gin"
)

// Cache 中间件：缓存 OpenAI 响应
func Cache(cache cache.Cache) gin.HandlerFunc {
	return func(c *gin.Context) {
		// 生成缓存 key（基于请求体）
		key := generateCacheKey(c)
		cached, err := cache.Get(key)
		if err == nil && len(cached) > 0 {
			c.Writer.Write(cached)
			c.Abort()
			return
		}
		// 创建缓冲区捕获响应体
		w := &bodyRewriteWriter{body: bytes.NewBufferString(""), ResponseWriter: c.Writer}
		// 继续执行后续逻辑
		c.Next()

		// 如果是成功响应，缓存结果
		if c.Writer.Status() == 200 {
			body := w.body.Bytes()
			cache.Set(key, body)
		}
	}
}

// 自定义 Writer 拦截响应体
type bodyRewriteWriter struct {
	gin.ResponseWriter
	body *bytes.Buffer
}

func (w *bodyRewriteWriter) Write(b []byte) (int, error) {
	w.body.Write(b) // 写入缓冲区
	return w.ResponseWriter.Write(b)
}

func generateCacheKey(c *gin.Context) string {
	// 简单实现：使用请求体的哈希作为 key
	body, _ := c.GetRawData()
	return fmt.Sprintf("cache:%x", body)
}
